#!/usr/bin/env python

import os
import sys
import types
import pygtk
pygtk.require('2.0')
import gtk
import gobject
import getdev
from multiprocessing import Process
import threading

class CustomSender(gobject.GObject):
    def __init__(self):
        self.__gobject_init__()

gobject.type_register(CustomSender)
gobject.signal_new('add_mi_finish', CustomSender, 
                    gobject.SIGNAL_RUN_FIRST,
                    #gobject.SIGNAL_RUN_LAST,
                    gobject.TYPE_NONE,
                    #(gobject.TYPE_PYOBJECT,)
                    ()
                )

class AddmiThread(threading.Thread):
    def __init__(self, usb_obj):
        threading.Thread.__init__(self)
        self.usb_obj = usb_obj
    def run(self):
        # Notice: we should emit "add_mi_finish" signal before return.
        paras = []
        if not self.usb_obj.get_parameters(paras):
            self.usb_obj.sender.emit('add_mi_finish')
            return
        (inst_usb_part, inst_usb_disk, iso_file) = paras
        usb_dev_path = self.usb_obj.usb_part_info[inst_usb_part]['devpath']
        usb_mount_dir = '/tmp/magicusbinst/%s' % inst_usb_part
        # Clear the msg box at first.
        self.usb_obj.showinfo('Install', '', clear=True)
        self.usb_obj.pbar.set_fraction(0.0)
        steps = 5
        svalue = 1.0/steps
        step = 0

        # mount the usb device.
        operate = 'Mount'
        self.usb_obj.pbar.set_text(operate)
        step = step + 1
        self.usb_obj.pbar.set_fraction(svalue * step)
        ret, msg = self.usb_obj.mount(usb_dev_path, usb_mount_dir)
        if not ret:
            self.usb_obj.showinfo(operate,
                     msg,
                     clear = False,
                     format = 'error')
            self.usb_obj.sender.emit('add_mi_finish')
            return

        # Copy grldr to usb.
        operate = 'Copy grldr'
        self.usb_obj.pbar.set_text(operate)
        step = step + 1
        self.usb_obj.pbar.set_fraction(svalue * step)
        ret, msg = self.usb_obj.copy_grldr(self.usb_obj.grldr, self.usb_obj.fonts, self.usb_obj.menu_lst, usb_mount_dir)
        if not ret:
            self.usb_obj.showinfo(operate,
                     msg,
                     clear = False,
                     format = 'warning')
        else:
            self.usb_obj.showinfo(operate,
                     'Successfully.\n',
                     clear = False,
                     format = 'success')

        # Extract MI kernel and rootfs from iso.
        operate = 'Extract MI file'
        self.usb_obj.pbar.set_text(operate)
        step = step + 1
        self.usb_obj.pbar.set_fraction(svalue * step)
        ret, msg = self.usb_obj.extract_mifiles(iso_file, usb_mount_dir)
        if not ret:
            self.usb_obj.showinfo(operate,
                     msg,
                     clear = False,
                     format = 'error')
            self.usb_obj.sender.emit('add_mi_finish')
            return
        else:
            self.usb_obj.showinfo(operate,
                     'Successfully.\n',
                     clear = False,
                     format = 'success')

        # umount the usb device.
        operate = 'Umount'
        self.usb_obj.pbar.set_text(operate)
        step = step + 1
        self.usb_obj.pbar.set_fraction(svalue * step)
        ret, msg = self.usb_obj.umount(usb_dev_path, usb_mount_dir)
        if not ret:
            self.usb_obj.showinfo(operate,
                     msg,
                     clear = False,
                     format = 'warning')
        # Remove the temp dir
        try:
            os.rmdir('/tmp/magicusbinst')
        except:
            pass

        operate = 'Write MBR'
        self.usb_obj.pbar.set_text(operate)
        step = step + 1
        self.usb_obj.pbar.set_fraction(svalue * step)
        ret, msg = self.usb_obj.write_usb_mbr(self.usb_obj.grldr_mbr, inst_usb_disk)
        if not ret:
            self.usb_obj.showinfo(operate,
                        msg,
                        clean = False,
                        format = 'error')

        self.usb_obj.pbar.set_text('Finish')
        self.usb_obj.sender.emit('add_mi_finish')
        return

class UsbInstaller():
    def __init__(self, usb_part_info, grldrt, grldr_mbr, fonts, menu_lst):

        self.usb_part_info = usb_part_info
        self.grldr = grldr
        self.grldr_mbr = grldr_mbr
        self.fonts = fonts
        self.menu_lst = menu_lst

        self.window = gtk.Window()
        self.window.connect('destroy', lambda w: gtk.main_quit())
        self.window.set_title('MagicUsbInstaller')
        self.window.set_border_width(20)
        #self.window.set_resizable(False)
        self.window.set_default_size(700, 300)
        self.main_hbox = gtk.HBox(True, 20)

        self.window.add(self.main_hbox)

        self.left_vbox = gtk.VBox(False, 10)
        self.right_vbox = gtk.VBox(False, 10)
        self.main_hbox.pack_start(self.left_vbox)
        self.main_hbox.pack_start(self.right_vbox)
        
        self.sw_msgbox = gtk.ScrolledWindow()
        self.sw_msgbox.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
        self.msgbox = gtk.TextView()
        self.msgbuffer = self.msgbox.get_buffer()
        self.msgbox.set_border_window_size(gtk.TEXT_WINDOW_LEFT, 30)
        #self.msgbox.set_border_window_size(gtk.TEXT_WINDOW_RIGHT, 30)
        self.msgbox.connect('expose_event', self.line_numbers_expose)
        self.sw_msgbox.add(self.msgbox)
        self.sw_msgbox.show()
        self.right_vbox.pack_start(self.sw_msgbox)
        self.msgbox.show()

        self.system_label = gtk.Label('Linux')
        self.system_label.set_alignment(0, 0)
        self.left_vbox.pack_start(self.system_label, False, False)
        self.system_label.show()

        self.part_combo = gtk.combo_box_new_text()
        self.left_vbox.pack_start(self.part_combo, False, False)
        self.part_combo.append_text('Select a partition:')
        for part in self.usb_part_info.keys():
            self.part_combo.append_text(part)
        self.part_combo.connect('changed', self.changed_cb)
        self.part_combo.set_active(0)
        self.part_combo.show()
       
        self.file_hbox = gtk.HBox(False, 3)
        self.file_text = gtk.Entry()
        self.file_text.set_text('/home/sunshine/magicinstaller2/result/MagicLinux-2.5-1.iso')
        self.file_hbox.pack_start(self.file_text)
        self.file_browse = gtk.Button('Browse')
        self.file_browse.connect('clicked', self.browse_file)
        self.file_hbox.pack_start(self.file_browse)
        self.file_browse.show()
        self.left_vbox.pack_start(self.file_hbox, False, False)
        self.file_text.show()
    
        self.add_del_hbox = gtk.HBox(False, 0)
        self.add_btn = gtk.Button('Install')
        #self.add_btn.set_tooltip_text('Installer the MI into usb devices')
        self.add_btn.connect('clicked', self.add_mi)
        self.del_btn = gtk.Button('Delete')
        self.del_btn.connect('clicked', self.del_mi)
        self.add_del_hbox.pack_start(self.add_btn)
        self.add_del_hbox.pack_start(self.del_btn)
        self.add_btn.show()
        self.del_btn.show()
        self.left_vbox.pack_start(self.add_del_hbox, False, False)
        
        self.pbar = gtk.ProgressBar()
        self.pbar_act = gtk.ProgressBar()
        self.left_vbox.pack_start(self.pbar, False, False)
        self.left_vbox.pack_start(self.pbar_act, False, False)
        self.pbar.show()
        self.pbar_act.show()

        self.left_vbox.show()
        self.right_vbox.show()
        self.main_hbox.show()
        
        self.window.show_all()
        return

    def get_parameters(self, result):
        inst_usb_part = ''
        inst_usb_disk = ''
        iso_file = self.file_text.get_text()

        model = self.part_combo.get_model()
        index = self.part_combo.get_active()
        if index:
            inst_usb_part = model[index][0]
            inst_usb_disk = self.usb_part_info[inst_usb_part]['disk']
        def check_para(title, value):
            pr_str = ''
            if not value.strip():
                pr_str = pr_str + 'none "%s" value, please check!' % title
            return pr_str
        pr_str = ''
        pr_str += '%s\n' % check_para('Usb Partition', inst_usb_part)
        pr_str += '%s\n' % check_para('Usb Disk Belong', inst_usb_disk)
        pr_str += '%s\n' % check_para('ISO File Path', iso_file)
        pr_str = pr_str.strip()
        if pr_str:
            self.showinfo('Can not get some parameters', pr_str)
            return False
        else:
            result.extend([inst_usb_part, inst_usb_disk, iso_file])
            return True

    def process_check(self):
        # For Process.
        #ret = self.add_process.is_alive()
        # For thread.
        ret = self.add_mi_thread.is_alive()
        #import pdb; pdb.set_trace()
        if ret == True:
            self.pbar_act.pulse()
            # As this is a timeout function, return TRUE so that it continues
            # to get called.
            return True
        else:
            self.pbar_act.set_fraction(0.0)
            self.time = 0
            return False

    def add_mi(self, widget):
        #self.add_process = Process(target=self.add_mi_process)
        #self.add_mi_process()

        self.close_button()
        self.sender = CustomSender()
        self.sender.connect('add_mi_finish', self.add_mi_cb)
        self.add_mi_thread = AddmiThread(self)
        #self.add_mi_thread.setDaemon(True)
        #self.add_mi_thread.daemon
        #import pdb; pdb.set_trace()
        self.timer = gobject.timeout_add(100, self.process_check)
        self.add_mi_thread.start()
        #self.add_mi_thread.join()
        #self.add_process.start()
        #self.timer = gobject.timeout_add(100, self.process_check)

    def add_mi_cb(self, widget):
        self.open_button()

    def close_button(self):
        self.add_btn.set_sensitive(False)
        self.del_btn.set_sensitive(False)

    def open_button(self):
        self.add_btn.set_sensitive(True)
        self.del_btn.set_sensitive(True)
        

    def write_usb_mbr(self, grldr_mbr, inst_usb_disk):
        msg = ''
        ret = True
        f_usb = None
        f_mbr = None
        new_mbr = ''
        if not inst_usb_disk:
            ret = False
            msg += 'The usb partition have none "devpath"\n parameter, please check!\n'
            return ret, msg

        # Construct the new mbr, keep the previous mbr's partition table.
        try:
            f_usb = open(inst_usb_disk, 'rb')
            f_mbr = open(grldr_mbr, 'rb')
            # We should keep the partition table.
            # The old version of grldr.mbr have 13 sector(6656 bytes), but the newest
            # grldr.mbr have 18 sector(9216 bytes).
            grub_mbr = f_mbr.read(9216)
            prev_mbr = f_usb.read(512)
            fs_beg = 0x1b8
            fs_len = 72
            addr = 0x0
            for s in grub_mbr:
                if addr < fs_beg or addr >= fs_beg + fs_len:
                    new_mbr += s
                elif addr >= fs_beg and addr < fs_beg + fs_len:
                    new_mbr += prev_mbr[addr]
                addr = addr+0x1
        except IOError, e:
            if f_mbr:
                f_mbr.close()
                f_mbr = None
            if f_usb:
                f_usb.close()
                f_usb = None
            ret = False
            msg += '%s\n' % str(e)
            return ret, msg
        f_mbr.close()
        f_mbr = None
        f_usb.close()
        f_usb = None

        # Write the new mbr, which we construct above, into usb device.
        try:
            f_usb = open(inst_usb_disk, 'wb')
            if new_mbr:
                f_usb.write(new_mbr)
        except IOError, e:
            if f_usb:
                f_usb.close()
                f_usb = None
            ret = False
            msg += '%s\n' % str(e)
            return ret, msg
        
        f_usb.close()
        f_usb = None
        return ret, msg

    def mount(self, dev, dir):
        msg = ''
        if not os.path.exists(dir):
            os.makedirs(dir)
        elif os.path.ismount(dir):
            # If it has been mounted, we should umount it firstly.
            msg += 'The dir %s is a mounted dir.\n' % dir
            cmd_res = self.run_bash('umount', [dir])
            ret = cmd_res['ret']
            if ret:
                msg += 'The return code is %s.\n' % str(ret)
                msg += '%s\n' % cmd_res['err']
                ret = not ret
                return ret, msg
        cmd_res = self.run_bash('mount', [dev, dir])
        ret = cmd_res['ret']
        if ret:
            msg += 'The return code is %s .\n' % str(ret)
            msg += '%s\n' % self.list_to_str(cmd_res['err'])
        ret = not ret
        return ret, msg

    def umount(self, dev, dir):
        msg = ''
        cmd_res = self.run_bash('umount', [dir])
        ret_mount = cmd_res['ret']
        if ret_mount:
            msg += 'The return code is %s.\n' % str(ret_mount)
            msg += '%s\n' % self.list_to_str(cmd_res['err'])
        ret_mount = not ret_mount

        ret_rmdir = True
        try:
            os.rmdir(dir)
        except OSError, e:
            ret_rmdir=False
            msg += 'Remove the temp dir "%s" failed.\n' % dir
            msg += '%s\n' % str(e)

        ret = ret_mount and ret_rmdir
        return ret, msg

    def copy_grldr(self, grldr, fonts, menu_lst, inst_usb_dir,
            mi_dir='magicinstaller'):
        msg = ''
        ret = True
        inst_dir = os.path.join(inst_usb_dir, mi_dir)
        
        if not os.path.exists(inst_usb_dir):
            try:
                os.mkdir(inst_usb_dir)
            except OSError,e:
                ret = False
                msg += 'Make dir "%s" failed.\n' % inst_usb_dir
                return ret, msg
        elif not os.path.exists(inst_dir):
            try:
                os.mkdir(inst_dir)
            except OSError, e:
                ret = False
                msg += 'Make dir "%s" failed.\n' % inst_dir
                return ret, msg

        def cp_file(file, inst_usb_dir):
            msg = ''
            cmd_res = self.run_bash('cp', ['-a', file, inst_usb_dir])
            cp_ret = cmd_res['ret']
            if cp_ret:
                msg += 'The return code is %s .\n' % str(cp_ret)
                msg += '%s\n' % self.list_to_str(cmd_res['err'])
            return cp_ret, msg

        ret_g, msg_g = cp_file(grldr, inst_usb_dir)
        msg += msg_g
        ret_f, msg_f = cp_file(fonts, os.path.join(inst_usb_dir, mi_dir))
        msg += msg_f
        ret_m, msg_m = cp_file(menu_lst, inst_usb_dir)
        msg += msg_m

        ret = not ret_g and not ret_f and not ret_m
        return ret, msg

    def extract_mifiles(self, iso_file, inst_usb_dir, mi_dir='magicinstaller'):
        msg = 'Extract mi file'
        inst_dir = os.path.join(inst_usb_dir, mi_dir)
        cmd_res = self.run_bash('7z', 
                            ['e', '-y', iso_file,
                             'boot/mirootfs.gz', 'boot/mivmlinuz',
                             '-o%s' % os.path.abspath(inst_dir)
                            ])
        ret = cmd_res['ret']
        if ret:
            msg += 'The return code is %s.\n' % ret
            msg += '%s\n' % cmd_res['err']
        ret = not ret
        return ret, msg

    def list_to_str(self, list):
        str = ''
        for l in list:
            str += '%s\n' % l.strip()
        return str.strip()

    def run_bash(self, cmd, argv, root='/'):
        import subprocess
        def chroot():
            os.chroot(root)

        cmd_res = {}
        res = subprocess.Popen([cmd] + argv, 
                                stdout = subprocess.PIPE, 
                                stderr = subprocess.PIPE, 
                                #preexec_fn = chroot, cwd = root,
                                close_fds = True)
        res.wait()
        cmd_res['out']=res.stdout.readlines()
        cmd_res['err']=res.stderr.readlines()
        cmd_res['ret']=res.returncode
        return cmd_res

    def del_mi(self, widget):
        pr_str = 'Plese delete MagicInstaller by yourself :P\n'
        pr_str+= 'You can remove the "grldr" file and the\n'
        pr_str+= '"magicinstaller" directory directly.\n'
        self.showinfo('Delete Way', pr_str)

    def browse_file(self, widget):
        dialog = gtk.FileChooserDialog( 'Open...', 
                                        None, 
                                        gtk.FILE_CHOOSER_ACTION_OPEN,
                                        (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
                                        gtk.STOCK_OPEN, gtk.RESPONSE_OK))
        dialog.set_default_response(gtk.RESPONSE_OK)
        filter = gtk.FileFilter()
        filter.set_name('ISO Files')
        filter.add_pattern('*.iso')
        dialog.add_filter(filter)
        
        filter = gtk.FileFilter()
        filter.set_name("All Files")
        filter.add_pattern("*")
        dialog.add_filter(filter)

        response = dialog.run()
        if response == gtk.RESPONSE_OK:
            #print dialog.get_filename(), 'selected'
            self.file_text.set_text(dialog.get_filename())
        elif response == gtk.RESPONSE_CANCEL:
            #print 'Closed, no files selected'
            pass
        dialog.destroy()
        

    def changed_cb(self, combobox):
        model = combobox.get_model()
        index = combobox.get_active()
        if index:
            usb_part = model[index][0]
            self.showinfo('USB INFO', self.usb_part_info[usb_part])
        return

    def get_lines(self, first_y, last_y, buffer_coords, numbers):
        text_view = self.msgbox
        iter, top = text_view.get_line_at_y(first_y)
        count = 0
        while not iter.is_end():
            y, height = text_view.get_line_yrange(iter)
            buffer_coords.append(y)
            line_num = iter.get_line()
            numbers.append(line_num)
            count += 1
            if(y + height) >= last_y:
                break
            iter.forward_line()
        return count

    def line_numbers_expose(self, widget, event, user_data=None):
        text_view = widget
        left_win = text_view.get_window(gtk.TEXT_WINDOW_LEFT)
        right_win = text_view.get_window(gtk.TEXT_WINDOW_RIGHT)
        if event.window == left_win:
            type = gtk.TEXT_WINDOW_LEFT
            target = left_win
        elif event.window == right_win:
            type = gtk.TEXT_WINDOW_RIGHT
            target = right_win
        else:
            return False
        first_y = event.area.y
        last_y = first_y + event.area.height

        #returns a tuple containing the textbuffer x and y coordinates
        #corresponding to the window_x and window_y coordinates in the
        #window specified by win
        x, first_y = text_view.window_to_buffer_coords(type, 0, first_y)
        x, last_y  = text_view.window_to_buffer_coords(type, 0, last_y)
        numbers = []
        pixels = []
        count = self.get_lines(first_y, last_y, pixels, numbers)
        layout = widget.create_pango_layout("")
        for i in range(count):
            x, pos = text_view.buffer_to_window_coords(type, 0, pixels[i])
            str = '%d' % numbers[i]
            layout.set_text(str)
            widget.style.paint_layout(target, widget.state, False,
                                      None, widget, None, 2, pos, layout)
        return False

    def showinfo(self, title, info, clear=True, format='normal'):
        #iter = textbuffer.get_start_iter()
        #iter = textbuffer.get_end_iter()
        #textbuffer.insert(iter, text)
        #textbuffer.set_text(text)
        pr_str = '%s:\n' % title
        if type(info) == types.DictType:
            keys = list(i for i in info.keys())
            keys.sort()
            for k in keys:
                pr_str = pr_str + '%s: %s\n' % (k, info[k])
        if type(info) == types.BooleanType:
            pass
        if type(info) == types.IntType:
            pass
        if type(info) == types.ListType:
            pass
        if type(info) == types.NoneType:
            pass
        if type(info) == types.StringType:
            pr_str = pr_str + info
        if type(info) == types.TupleType:
            pass
        if clear:
            self.msgbuffer.set_text(pr_str.decode('gbk'))
        else:
            iter = self.msgbuffer.get_end_iter()
            self.msgbuffer.insert(iter, pr_str.decode('gbk'))

    def main(self):
        gtk.main()
        return



if 1:
    def get_usb_part(usb_map, all_part_info, disk_size_map):
        usb_part_info = {}
        def get_attr(tmap, attr):
            if not tmap.has_key(attr):
                return None
            else:
                return tmap[attr]

        for disk in usb_map.keys():
            for part in [disk] + usb_map[disk]:
                if all_part_info[part].has_key('ID_FS_TYPE'):
                    # Disk maybe a partition... , so we check the disk too.
                    # disk:     partition belong to this disk.
                    # devpath:  the partition's dev path.
                    # vendor:   the vendor of the usb.
                    # model:    the flash model of the usb.
                    # model_id: the flash model's id of the usb.
                    # fstype:   the filesystem type of this partition.
                    # fsver:    the filesystem's version of this partition.
                    # size:     the size of this partition.
                    part_info = {}
                    part_info['disk'] = get_attr(all_part_info[disk], 'DEVNAME')
                    part_info['devpath'] = get_attr(all_part_info[part],'DEVNAME')
                    part_info['vendor'] = get_attr(all_part_info[part],'ID_VENDOR')
                    part_info['model'] = get_attr(all_part_info[part],'ID_MODEL')
                    part_info['model_id'] = get_attr(all_part_info[part],'ID_MODEL_ID')
                    part_info['fstype'] = get_attr(all_part_info[part],'ID_FS_TYPE')
                    part_info['fsver'] = get_attr(all_part_info[part],'ID_FS_VERSION')
                    part_info['size'] = '%.2f MB' % (float(disk_size_map[part])/1024/1024)
                    usb_part_info[part] = part_info
        return usb_part_info

if __name__ == "__main__":
    devices_info, devices_size = getdev.probe(getdev.ALL_DEVICE_INFO)
    usb_map = getdev.probe(getdev.CLASS_USB)
    usb_part_info = get_usb_part( usb_map,
                                  devices_info,
                                  devices_size )
    grldr = os.path.abspath('grub/grldr')
    grldr_mbr = os.path.abspath('grub/grldr.mbr')
    fonts = os.path.abspath('grub/fonts')
    menu_lst = os.path.abspath('grub/menu.lst')
    usbInst = UsbInstaller( usb_part_info, grldr, grldr_mbr, fonts, menu_lst )
    usbInst.main()

